from queue import Queue
import random
import threading

# utils function
def is_in_network(ip, network):
    # 简单示例，假设网络格式为 'xxx.xxx.xxx.0/24'
    network_ip, prefix = network.split('/')
    network_parts = list(map(int, network_ip.split('.')))
    ip_parts = list(map(int, ip.split('.')))
    mask = (0xFFFFFFFF << (32 - int(prefix))) & 0xFFFFFFFF
    network_int = (network_parts[0] << 24) | (network_parts[1] << 16) | (network_parts[2] << 8) | network_parts[3]
    ip_int = (ip_parts[0] << 24) | (ip_parts[1] << 16) | (ip_parts[2] << 8) | ip_parts[3]
    return (ip_int & mask) == (network_int & mask)

def get_network_from_ip(ip_mask):
    ip, prefix = ip_mask.split('/')
    ip_parts = list(map(int, ip.split('.')))
    mask = (0xFFFFFFFF << (32 - int(prefix))) & 0xFFFFFFFF
    network_int = (ip_parts[0] << 24) | (ip_parts[1] << 16) | (ip_parts[2] << 8) | ip_parts[3]
    network_ip = (network_int & mask)
    network_parts = [
        (network_ip >> 24) & 0xFF,
        (network_ip >> 16) & 0xFF,
        (network_ip >> 8) & 0xFF,
        network_ip & 0xFF
    ]
    return f"{network_parts[0]}.{network_parts[1]}.{network_parts[2]}.{network_parts[3]}/{prefix}"

# global
running = True
port_rcv_fifo = Queue()

# 定义 EthernetFrame 类，包含二层信息和上层协议类型
class EthernetFrame:
    def __init__(self, src_mac, dst_mac, eth_type, payload):
        self.src_mac = src_mac
        self.dst_mac = dst_mac
        self.eth_type = eth_type  # 上层协议类型，例如 'IPv4', 'ARP'
        self.payload = payload  # 上层协议数据（例如 IPv4Packet 或 ARPMessage 实例）

    def __str__(self):
        return f"Ether {self.src_mac} > {self.dst_mac}|{self.payload}"

# 定义 IPv4Packet 类，包含三层信息
class IPv4Packet:
    def __init__(self, src_ip=None, dst_ip=None, protocol=None, ttl=64, payload=None):
        self.src_ip = src_ip
        self.dst_ip = dst_ip
        self.protocol = protocol
        self.ttl = ttl
        self.payload = payload

    def __str__(self):
        return f"IP {self.src_ip} > {self.dst_ip}, ttl {self.ttl}|{self.payload}"

# 定义 ICMPMessage 类，包含四层信息
class ICMPMessage:
    def __init__(self, icmp_type, icmp_id, data, sequence=1):
        self.icmp_type = icmp_type
        self.icmp_id = icmp_id
        self.sequence = sequence
        self.data = data

    def __str__(self):
        return f"{self.icmp_type} id {self.icmp_id}, seq {self.sequence}, data {self.data}"

# 定义 ARPMessage 类，包含 ARP 协议信息
class ARPMessage:
    def __init__(self, arp_type, src_mac, src_ip, dst_mac, dst_ip):
        self.arp_type = arp_type  # 'ARP Request' 或 'ARP Response'
        self.src_mac = src_mac
        self.src_ip = src_ip
        self.dst_mac = dst_mac
        self.dst_ip = dst_ip

    def __str__(self):
        return f"{self.arp_type} {self.src_ip} > {self.dst_ip}"

# 定义 Port 类
class Port:
    def __init__(self, device, index, mac_address=None, ip_mask=None):
        self.device = device
        self.index = index
        self.mac_address = mac_address
        self.ip_address = None
        if ip_mask:
            self.ip_address, self.prefix = ip_mask.split('/')
            self.network = get_network_from_ip(ip_mask)
        self.link = None
        # 等待发送的数据包列表{arp_dst_ip: [pkt1, pkt2]}（比如未获取目的mac的数据包，需要先缓存，等arp响应到达之后再发送）
        self.waiting_packets = {}

    def connect_link(self, link):
        self.link = link

# 定义一个抽象的网络设备类
class NetworkDevice:
    def __init__(self, name, mac_address=None, ip_mask=None):
        self.name = name
        # 接口列表
        self.ports = {}
        if mac_address and ip_mask:
            self.add_port('eth1', mac_address, ip_mask)

    def add_port(self, index, mac_address=None, ip_mask=None):
        port = Port(self, index, mac_address, ip_mask)
        self.ports[index] = port
        return port

    def add_route(self, network, next_hop, route_type='静态'):
        print(f'[{self.name}] add route {network} -> {next_hop} ({route_type})')
        self.routing_table[network] = (next_hop, route_type)

    def send_packet(self, packet, out_port_index=None):
        print('no implement send_packet')

    def get_route_for_ip(self, ip):
        print(f'[{self.name}] routing table: {self.routing_table}')
        for network, (next_hop, route_type) in self.routing_table.items():
            if is_in_network(ip, network):
                return next_hop, route_type
        return None, None

# 定义交换机类，继承自 NetworkDevice
class Switch(NetworkDevice):
    def __init__(self, name):
        super().__init__(name)
        self.mac_table = {}

    def receive_packet(self, packet, in_port):
        # L2 switch logic
        # 1. 学习源 MAC 地址
        self.mac_table[packet.src_mac] = in_port.index
        print(f"交换机 {self.name} 学习到 MAC 地址 {packet.src_mac} 位于端口 {in_port.index}")

        # 2. 如果目的 MAC 地址在 MAC 地址表中，将数据包转发到对应的端口
        if packet.dst_mac in self.mac_table:
            target_port_index = self.mac_table[packet.dst_mac]
            target_port = self.ports.get(target_port_index)
            print(f"交换机 {self.name} 将数据包从端口 {in_port.index} 转发到端口 {target_port_index}")
            target_port.link.transmit(packet, target_port)
        # 3. 如果目的 MAC 地址不在 MAC 地址表中，将数据包广播到所有端口
        else:
            print(f"交换机 {self.name} 未找到 MAC 地址 {packet.dst_mac}，进行广播")
            for port_index, port in self.ports.items():
                # 跳过输入端口(避免环路)，和未连接的端口
                if port != in_port and port.link:
                    print(f'{self.name} send packet to port {port_index}')
                    port.link.transmit(packet, port)

    def send_packet(self, packet, out_port_index):
        print('switch itself should not send packet')

# 定义主机类，继承自 NetworkDevice
class Host(NetworkDevice):
    def __init__(self, name, mac_address, ip_mask):
        super().__init__(name, mac_address, ip_mask)
        # 默认路由下一跳
        self.default_gateway = None
        # ARP 表，格式：{IP 地址: MAC 地址}
        self.arp_table = {}
        self.mac_address = mac_address
        self.ip_address = ip_mask.split('/')[0] if ip_mask else None
        self.routing_table = {}
        self.icmp_requests = {}  # 用于存储 ICMP 请求的字典，key 为 icmp_id，value 为请求信息
        if ip_mask:
            # 添加直连路由
            port = self.ports['eth1']
            self.add_route(port.network, port.ip_address, '直连')

    def set_default_gateway(self, gateway_ip):
        self.default_gateway = gateway_ip

    def add_route(self, network, next_hop, route_type='静态'):
        super().add_route(network, next_hop, route_type)

    def ping(self, ip, count=1):
        print(f'[{self.name}]# ping {ip} from {self.ip_address}')
        icmp_id = random.randint(1, 65535)  # 生成一个随机的 icmp_id
        for i in range(count):
            icmp_message = ICMPMessage(icmp_type='ICMP Request', icmp_id=icmp_id, data="Hello, Friend!", sequence=i+1)
            ipv4_packet = IPv4Packet(src_ip=self.ip_address, dst_ip=ip, payload=icmp_message)
            ethernet_frame = EthernetFrame(src_mac=self.mac_address, dst_mac="FF:FF:FF:FF:FF:FF", eth_type='IPv4', payload=ipv4_packet)
            self.icmp_requests[icmp_id] = icmp_message  # 存储 ICMP 请求
            self.send_packet(ethernet_frame)

    def send_packet(self, packet):
        default_port = self.ports['eth1']
        # 1. 查找路由表
        next_hop, route_type = self.get_route_for_ip(packet.payload.dst_ip)
        print(f'lookup route for {packet.payload.dst_ip}: route_type: {route_type}')
        if route_type == '直连':
            print(f'主机 {self.ip_address} 目的IP {packet.payload.dst_ip} 在直连网段，查找目的主机ARP')
            arp_dst_ip = packet.payload.dst_ip
        elif route_type == '静态':
            print(f'主机 {self.ip_address} 开始查找路由下一跳ARP')
            arp_dst_ip = next_hop
        else:
            print(f"通过默认网关({self.default_gateway}) 发送，查找默认网关ARP")
            arp_dst_ip = self.default_gateway
        # 查看ARP表
        dst_mac = self.arp_table.get(arp_dst_ip)
        if not dst_mac:
            print(f"主机 {self.ip_address} ARP 表中没有 {arp_dst_ip} 的 MAC 地址，发送 ARP 请求")
            # 缓存IP报文
            if arp_dst_ip not in default_port.waiting_packets:
                default_port.waiting_packets[arp_dst_ip] = []
            default_port.waiting_packets[arp_dst_ip].append(packet)
            arp_message = ARPMessage(arp_type="ARP Request", src_mac=self.mac_address, src_ip=self.ip_address, dst_mac="00:00:00:00:00:00", dst_ip=arp_dst_ip)
            arp_ethernet_frame = EthernetFrame(src_mac=self.mac_address, dst_mac="FF:FF:FF:FF:FF:FF", eth_type='ARP', payload=arp_message)
            if default_port.link:
                default_port.link.transmit(arp_ethernet_frame, default_port)
        else:
            packet.dst_mac = dst_mac
            print(f"主机 {self.ip_address} 查找ARP表，将数据包发往（MAC: {dst_mac}）")
            default_port.link.transmit(packet, default_port)

    def receive_packet(self, packet, in_port):
        if packet.eth_type == 'ARP':
            arp_message = packet.payload
            if arp_message.dst_ip != in_port.ip_address:
                print(f"主机 {self.ip_address} 收到非本机 ARP 请求，drop")
                return
            if arp_message.arp_type == "ARP Request":
                print(f"主机 {self.ip_address} 收到 ARP 请求，发送 ARP 响应")
                arp_response = ARPMessage(arp_type="ARP Response", src_mac=self.mac_address, src_ip=self.ip_address, dst_mac=arp_message.src_mac, dst_ip=arp_message.src_ip)
                arp_response_ethernet_frame = EthernetFrame(src_mac=self.mac_address, dst_mac=arp_message.src_mac, eth_type='ARP', payload=arp_response)
                in_port.link.transmit(arp_response_ethernet_frame, in_port)
            elif arp_message.arp_type == "ARP Response":
                print(f"主机 {self.ip_address} 收到 ARP 响应，更新 ARP 表")
                self.arp_table[arp_message.src_ip] = arp_message.src_mac
                # 如果有等待此arp的IP报文，则发送
                for waiting_packet in in_port.waiting_packets.get(arp_message.src_ip, []):
                    print(f"主机 {self.ip_address} 发送等待的 {waiting_packet.payload.payload.icmp_type} 数据包")
                    waiting_packet.dst_mac = arp_message.src_mac
                    in_port.link.transmit(waiting_packet, in_port)
                # 发送完清楚包缓存
                del in_port.waiting_packets[arp_message.src_ip]
        elif packet.eth_type == 'IPv4':
            ipv4_packet = packet.payload
            if ipv4_packet.dst_ip == self.ip_address:
                if isinstance(ipv4_packet.payload, ICMPMessage):
                    icmp_message = ipv4_packet.payload
                    if icmp_message.icmp_type == "ICMP Request":
                        print(f"主机 {self.ip_address} 收到 ICMP 请求，发送 ICMP 响应")
                        icmp_response = ICMPMessage(icmp_type='ICMP Response', icmp_id=icmp_message.icmp_id, data=icmp_message.data)
                        ipv4_response = IPv4Packet(src_ip=self.ip_address, dst_ip=ipv4_packet.src_ip, payload=icmp_response)
                        ethernet_response = EthernetFrame(src_mac=self.mac_address, dst_mac=packet.src_mac, eth_type='IPv4', payload=ipv4_response)
                        in_port.link.transmit(ethernet_response, in_port)
                    elif icmp_message.icmp_type == "ICMP Response":
                        print(f"主机 {self.ip_address} 收到 ICMP 响应")
                        request_packet = self.icmp_requests.get(icmp_message.icmp_id)
                        if request_packet:
                            print(f"主机 {self.ip_address} 匹配到 ICMP 请求，请求数据: {request_packet.data}")
                            del self.icmp_requests[icmp_message.icmp_id]  # 删除已处理的请求
                        else:
                            print(f"主机 {self.ip_address} 未找到匹配的 ICMP 请求")
                else:
                    print(f"主机 {self.ip_address} 接收到来自 {ipv4_packet.src_ip} 的数据包，自定义数据: {ipv4_packet.payload.data}")
                    # 处理数据包
            else:
                print(f"主机 {self.ip_address} 接收到非目的数据包，丢弃")

# 定义路由器类，继承自 NetworkDevice
class Router(NetworkDevice):
    def __init__(self, name, mac_address, ip_with_prefix):
        super().__init__(name, mac_address, ip_with_prefix)
        # 路由器的路由表，格式：{网段: (下一跳 IP, 类型)}
        self.routing_table = {}
        # ARP 表，格式：{IP 地址: MAC 地址}
        self.arp_table = {}
        if ip_with_prefix:
            # 添加直连路由
            port = self.ports['eth1']
            self.add_route(port.network, port.ip_address, '直连')

    def add_interface(self, index, mac_address, ip_mask):
        port = self.add_port(index, mac_address, ip_mask)
        self.add_route(port.network, port.ip_address, '直连')
        return port

    def send_packet(self, packet, in_port):
        # 1. 查找路由表
        next_hop, route_type = self.get_route_for_ip(packet.payload.dst_ip)
        out_port = None
        print(f'[{self.name}] 1. lookup route for {packet.payload.dst_ip}: next_hop: {next_hop}')
        # 2. 根据路由类型找out_port
        if not next_hop:
            print(f"路由器 {self.name} 没有找到合适的路由，丢弃数据包")
            # todo: reply icmp message(network unreachable)
        if route_type == '直连':
            print(f"路由器 {self.name} 发现直连网段，直接发送数据包")
            out_port = [port for port in self.ports.values() if port.ip_address == next_hop][0]
            arp_dst_ip = packet.payload.dst_ip
        elif route_type == '静态':
            print(f"路由器 {self.name} 发现静态路由，根据下一跳IP({next_hop})找出接口")
            #
            out_port = [port for port in self.ports.values() if is_in_network(next_hop, port.network)]
            if not out_port:
                print(f"路由器 {self.name} 下一跳network不可达，丢弃数据包")
                return
            arp_dst_ip = next_hop
        # 3. 查找下一跳IP的ARP表
        packet.src_mac = out_port.mac_address
        dst_mac = self.arp_table.get(arp_dst_ip)
        if not dst_mac:
            # cache pkt
            if arp_dst_ip not in out_port.waiting_packets:
                out_port.waiting_packets[arp_dst_ip] = []
            out_port.waiting_packets[arp_dst_ip].append(packet)
            print(f"路由器 {self.name} ARP 表中没有 {arp_dst_ip} 的 MAC 地址，发送 ARP 请求")
            arp_message = ARPMessage(arp_type="ARP Request", src_mac=out_port.mac_address, src_ip=out_port.ip_address, dst_mac="00:00:00:00:00:00", dst_ip=arp_dst_ip)
            arp_ethernet_frame = EthernetFrame(src_mac=out_port.mac_address, dst_mac="FF:FF:FF:FF:FF:FF", eth_type='ARP', payload=arp_message)
            out_port.link.transmit(arp_ethernet_frame, out_port)
        else:
            packet.dst_mac = dst_mac
            print(f"路由器 {self.name} 查找路由表，将数据包发往 {arp_dst_ip}（MAC: {dst_mac}）")
            out_port.link.transmit(packet, out_port)

    def receive_packet(self, packet, in_port):
        print(f"[{self.name}]:{in_port.index} recv {packet.eth_type}：[{packet.payload.src_ip}->{packet.payload.dst_ip}]")
        # 1. 处理local packet
        if packet.payload.dst_ip == in_port.ip_address:
            if packet.eth_type == "ARP":
                arp_message = packet.payload
                if arp_message.arp_type == "ARP Request":
                    print(f"[{self.name}] 收到 ARP 请求，发送 ARP 响应")
                    arp_response = ARPMessage(arp_type="ARP Response",
                                              src_mac=in_port.mac_address,
                                              src_ip=arp_message.dst_ip,
                                              dst_mac=arp_message.src_mac,
                                              dst_ip=arp_message.src_ip)
                    arp_response_ethernet_frame = EthernetFrame(src_mac=in_port.mac_address, dst_mac=arp_message.src_mac, eth_type='ARP', payload=arp_response)
                    in_port.link.transmit(arp_response_ethernet_frame, in_port)
                elif arp_message.arp_type == "ARP Response":
                    print(f"[{self.name}] 收到 ARP 响应，更新 ARP 表")
                    self.arp_table[arp_message.src_ip] = arp_message.src_mac
                    # 如果有等待此arp的IP报文，则发送
                    for waiting_packet in in_port.waiting_packets.get(arp_message.src_ip, []):
                        print(f"[{self.name}]:{in_port.index} 发送等待的 {waiting_packet.payload.payload.icmp_type} 数据包")
                        waiting_packet.dst_mac = arp_message.src_mac
                        in_port.link.transmit(waiting_packet, in_port)
            elif packet.eth_type == "IPv4":
                ipv4_packet = packet.payload
                print(f"[{self.name}] 接收到来自 {ipv4_packet.src_ip} 的数据包，自定义数据: {ipv4_packet.payload.data}")
                # todo: icmp reply
        else:
            if packet.eth_type == "ARP":
                print(f"[{self.name}] 收到非本设备 ARP 请求，drop")
                return
            # 非本机报文，通过路由转发
            self.send_packet(packet, in_port)


# 定义链路类
class Link:
    def __init__(self, port1, port2):
        self.port1 = port1
        self.port2 = port2
        port1.connect_link(self)
        port2.connect_link(self)

    def transmit(self, packet, in_port):
        if in_port not in [self.port1, self.port2]:
            print('链路错误: 端口不在链路内')
        if in_port == self.port1:
            target_port = self.port2
        else:
            target_port = self.port1
        print(f'port_tx: [{in_port.device.name}]:{in_port.index} {packet}]')
        port_rcv_fifo.put((packet, target_port))

# 网络线路收发模拟器(设备需要发包时，只需发送到对应Link实体，自动会触发Link另一端设备的收包回调)
# 这里通过一个fifo实现发送端与接收端之间的解耦（如果在发送端直接调用对端的接收，可能会造成函数栈过深）
# 效果类似于设备网卡收包中断后调用收包回调函数。理论上每个设备创建一个fifo及收包线程更符合实际情况，这里统一用
# 一个全局fifo和收包线程，通过记录port与pkt对应关系来处理是为了更高效，相当于io多路复用。

def global_port_rcv_loop():
    global running
    print('[thread][global_port_rcv_loop] start')
    # 处理完再退出
    while running or port_rcv_fifo.qsize() > 0:
        try:
            packet, in_port = port_rcv_fifo.get()
            print(f'port_rx: [{in_port.device.name}]:{in_port.index} {packet}]')
            in_port.device.receive_packet(packet, in_port)
        except Exception as e:
            print(f'Error in global_port_rcv_loop: {e}')
    print('[thread][global_port_rcv_loop] exit')


if __name__ == '__main__':
    # 创建线路接收处理线程
    thread = threading.Thread(target=global_port_rcv_loop, daemon=True)
    thread.start()
    # 创建设备实例
    host1 = Host('host1', "00:11:22:33:44:11", "192.168.1.100/24")
    host2 = Host('host2', "00:11:22:33:44:22", "192.168.1.200/24")
    host3 = Host('host3', "00:11:22:33:44:33", "192.168.2.100/24")
    switch = Switch('sw1')
    router = Router('router1', "AA:22:33:44:55:11", "192.168.1.254/24")

    # 添加端口和接口
    switch_port1 = switch.add_port(1)
    switch_port2 = switch.add_port(2)
    switch_port3 = switch.add_port(3)
    router_port2 = router.add_interface('eth2', "AA:22:33:44:55:22", "192.168.2.254/24")

    # 创建链路实例
    link1 = Link(switch_port1, host1.ports['eth1'])
    link2 = Link(switch_port2, host2.ports['eth1'])
    link3 = Link(switch_port3, router.ports['eth1'])
    link4 = Link(router_port2, host3.ports['eth1'])

    # 配置主机 1 的路由表和默认网关
    host1.set_default_gateway("192.168.1.254")
    host2.set_default_gateway("192.168.2.254")

    # 发包测试
    print("==========同网段测试：host1 ping host2...")
    #host1.ping(host2.ip_address, count=1)
    # print("==========不同网段测试：host1 ping host3...")
    host1.ping(host3.ip_address, count=1)
    running = False
    thread.join()